import os
import re

def clean_processed_directory(processed_dir: str):
    """
    清理processed目录，删除旧格式的数据文件。
    
    旧格式: sample_XXXXX.npy
    新格式: XXX-X_sample_X.npy (例如: 001-1_sample_0.npy)
    """
    # 确保目录存在
    if not os.path.exists(processed_dir):
        print(f"目录不存在: {processed_dir}")
        return
        
    # 旧格式文件的正则模式
    old_pattern = re.compile(r'^sample_\d+\.npy$')
    # 新格式文件的正则模式
    new_pattern = re.compile(r'^\d{3}-\d+_sample_\d+\.npy$')
    
    # 统计计数器
    old_files = []
    new_files = []
    
    # 遍历目录
    for filename in os.listdir(processed_dir):
        if old_pattern.match(filename):
            old_files.append(filename)
        elif new_pattern.match(filename):
            new_files.append(filename)
    
    print(f"找到 {len(old_files)} 个旧格式文件")
    print(f"找到 {len(new_files)} 个新格式文件")
    
    # 确认是否删除
    if old_files:
        confirmation = input(f"确认要删除 {len(old_files)} 个旧格式文件吗? (y/n): ")
        if confirmation.lower() == 'y':
            # 删除旧文件
            for file in old_files:
                file_path = os.path.join(processed_dir, file)
                try:
                    os.remove(file_path)
                    print(f"已删除: {file}")
                except Exception as e:
                    print(f"删除文件 {file} 时出错: {e}")
            print("清理完成！")
        else:
            print("操作已取消")
    else:
        print("没有找到需要删除的旧格式文件")

if __name__ == '__main__':
    # 处理训练集和测试集
    data_dirs = {
        'train_data': 'data/processed/train_data',
        'test_data': 'data/processed/test_data'
    }
    
    for name, dir_path in data_dirs.items():
        print(f"\n处理{name}...")
        clean_processed_directory(dir_path)